#include "cost_layer.h"
#include "owl_math.h"
#include "logging.h"

CostLayer MakeCostLayer(int batch, int inputs, int type, Float scale) {
    CostLayer layer = {0};
    layer.type = COST;
    layer.scale = scale;
    layer.batch = batch;
    layer.inputs = inputs;
    layer.outputs = inputs;
    layer.data_delta = (Float *)OwlAlloc(inputs * batch , sizeof(Float));
    layer.data_output = (Float *)OwlAlloc(inputs * batch, sizeof(Float));
    layer.cost = (Float *)OwlAlloc(1, sizeof(Float));

    layer.Forward = ForwardCostLayer;
    layer.Backward = BackwardCostLayer;
    return layer;
}

void ForwardCostLayer(const CostLayer layer, NetworkState state) {
    LOGFUNC("ForwardCostLayer");
    if (!state.data_truth) {
        return;
    }
    OWLLOG("ForwardCostLayer input", state.data_input, 3);
    OwlL2(layer.inputs * layer.batch, state.data_input, state.data_truth, layer.data_delta, layer.data_output);
    OWLLOG("ForwardCostLayer", layer.data_delta, layer.batch * layer.inputs);
    layer.cost[0] = OwlSum(layer.batch * layer.outputs, layer.data_output);
}

void BackwardCostLayer(const CostLayer layer, NetworkState state) {
    LOG_FUNC();//("BackwardCostLayer");
    OwlAxpy(layer.batch * layer.inputs, layer.scale, layer.data_delta, 1, state.data_delta, 1);
    OWLLOG("BackwardCostLayer", state.data_delta, layer.batch * layer.inputs);
}
